#############################################################   版本一：需要手动输入运行指令 + 输入任务描述 #############################################################################
""" 
python -m lerobot.wwq_record \
    --robot.type=so100_follower \
    --robot.port=/dev/ttyACM14 \
    --robot.cameras="{top: {type: opencv, index_or_path: 2, width: 640, height: 480, fps: 30}, laptop: {type: opencv, index_or_path: 4, width: 640, height: 480, fps: 30}}" \
    --robot.id=my_awesome_follower_arm \
    --display_data=false \
    --dataset.repo_id=datasets/eval_wwq_pencil \
    --policy.path=/home/zhq2004/XArm/lerobot_smolvla/outputs/train/smolvla_so100_pencil/checkpoints/160000/pretrained_model \
    --dataset.push_to_hub=False \
    --dataset.num_episodes=1 \
    --dataset.episode_time_s=30 \
    --dataset.root=/home/zhq2004/XArm/lerobot_smolvla/datasets/eval_pick_place_red_smolvla2


请输入任务描述（例如：'Place the black pen inside the pen holder'）：
Place the black pen inside the pen holder
"""

# import logging
# import time
# from dataclasses import asdict, dataclass
# from pathlib import Path
# from pprint import pformat
# from typing import List

# from lerobot.cameras import (  # noqa: F401
#     CameraConfig,  # noqa: F401
# )
# from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
# from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
# from lerobot.configs import parser
# from lerobot.configs.policies import PreTrainedConfig
# from lerobot.datasets.image_writer import safe_stop_image_writer
# from lerobot.datasets.lerobot_dataset import LeRobotDataset
# from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
# from lerobot.policies.factory import make_policy
# from lerobot.policies.pretrained import PreTrainedPolicy
# from lerobot.robots import (  # noqa: F401
#     Robot,
#     RobotConfig,
#     koch_follower,
#     make_robot_from_config,
#     so100_follower,
#     so101_follower,
# )
# from lerobot.teleoperators import (  # noqa: F401
#     Teleoperator,
#     TeleoperatorConfig,
#     koch_leader,
#     make_teleoperator_from_config,
#     so100_leader,
#     so101_leader,
# )
# from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop
# from lerobot.utils.control_utils import (
#     init_keyboard_listener,
#     is_headless,
#     predict_action,
#     sanity_check_dataset_name,
#     sanity_check_dataset_robot_compatibility,
# )
# from lerobot.utils.robot_utils import busy_wait
# from lerobot.utils.utils import (
#     get_safe_torch_device,
#     init_logging,
#     log_say,
# )
# from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data


# """ 
# record.py 是 LeRobot 框架中用于录制机器人数据集的脚本，主要功能包括：

#     数据采集：从机器人（例如 so100_follower）和摄像头（例如 OpenCV 摄像头）获取观测数据（图像、状态）。
#     动作生成：通过远程操作（teleop）或策略（policy）生成机器人动作。
#     数据集保存：将观测和动作保存为数据集，存储到本地（--dataset.root）或推送到 Hugging Face Hub（--dataset.push_to_hub）。
#     实时控制：以指定帧率（--dataset.fps）控制机器人动作和数据录制。
# """

# # 定义数据集录制的参数
# @dataclass
# class DatasetRecordConfig:
#     # Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).
#     repo_id: str
#     # A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.")
#     single_task: str | None = None  # 修改为可选字段
#     # Root directory where the dataset will be stored (e.g. 'dataset/path').
#     root: str | Path | None = None
#     # Limit the frames per second.
#     fps: int = 30
#     # Number of seconds for data recording for each episode.
#     episode_time_s: int | float = 60
#     # Number of seconds for resetting the environment after each episode.
#     reset_time_s: int | float = 60
#     # Number of episodes to record.
#     num_episodes: int = 50
#     # Encode frames in the dataset into video
#     video: bool = True
#     # Upload dataset to Hugging Face hub.
#     push_to_hub: bool = True
#     # Upload on private repository on the Hugging Face hub.
#     private: bool = False
#     # Add tags to your dataset on the hub.
#     tags: list[str] | None = None
#     # Number of subprocesses handling the saving of frames as PNG. Set to 0 to use threads only;
#     # set to ≥1 to use subprocesses, each using threads to write images. The best number of processes
#     # and threads depends on your system. We recommend 4 threads per camera with 0 processes.
#     # If fps is unstable, adjust the thread count. If still unstable, try using 1 or more subprocesses.
#     num_image_writer_processes: int = 0
#     # Number of threads writing the frames as png images on disk, per camera.
#     # Too many threads might cause unstable teleoperation fps due to main thread being blocked.
#     # Not enough threads might cause low camera fps.
#     num_image_writer_threads_per_camera: int = 4

# # 定义整体录制配置，包括机器人、数据集和控制方式
# @dataclass
# class RecordConfig:
#     robot: RobotConfig
#     dataset: DatasetRecordConfig
#     # Whether to control the robot with a teleoperator
#     teleop: TeleoperatorConfig | None = None
#     # Whether to control the robot with a policy
#     policy: PreTrainedConfig | None = None
#     # Display all cameras on screen
#     display_data: bool = False
#     # Use vocal synthesis to read events.
#     play_sounds: bool = True
#     # Resume recording on an existing dataset.
#     resume: bool = False

#     def __post_init__(self):
#         # HACK: We parse again the cli args here to get the pretrained path if there was one.
#         policy_path = parser.get_path_arg("policy")
#         if policy_path:
#             cli_overrides = parser.get_cli_overrides("policy")
#             self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
#             self.policy.pretrained_path = policy_path

#         if self.teleop is None and self.policy is None:
#             raise ValueError("Choose a policy, a teleoperator or both to control the robot")

#     @classmethod
#     def __get_path_fields__(cls) -> list[str]:
#         """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
#         return ["policy"]


# @safe_stop_image_writer
# def record_loop(
#     robot: Robot,
#     events: dict,
#     fps: int,
#     dataset: LeRobotDataset | None = None,
#     teleop: Teleoperator | List[Teleoperator] | None = None,
#     policy: PreTrainedPolicy | None = None,
#     control_time_s: int | None = None,
#     single_task: str | None = None,
#     display_data: bool = False,
# ):
#     if dataset is not None and dataset.fps != fps:
#         raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).")

#     teleop_arm = teleop_keyboard = None
#     if isinstance(teleop, list):
#         teleop_keyboard = next((t for t in teleop if isinstance(t, KeyboardTeleop)), None)
#         teleop_arm = next(
#             (
#                 t
#                 for t in teleop
#                 if isinstance(t, (so100_leader.SO100Leader, so101_leader.SO101Leader, koch_leader.KochLeader))
#             ),
#             None,
#         )

#         if not (teleop_arm and teleop_keyboard and len(teleop) == 2 and robot.name == "lekiwi_client"):
#             raise ValueError(
#                 "For multi-teleop, the list must contain exactly one KeyboardTeleop and one arm teleoperator. Currently only supported for LeKiwi robot."
#             )

#     # if policy is given it needs cleaning up
#     if policy is not None:
#         policy.reset()

#     timestamp = 0
#     start_episode_t = time.perf_counter()
#     while timestamp < control_time_s:
#         start_loop_t = time.perf_counter()

#         if events["exit_early"]:
#             events["exit_early"] = False
#             break

#         observation = robot.get_observation()

#         if policy is not None or dataset is not None:
#             observation_frame = build_dataset_frame(dataset.features, observation, prefix="observation")

#         if policy is not None:
#             action_values = predict_action(
#                 observation_frame,
#                 policy,
#                 get_safe_torch_device(policy.config.device),
#                 policy.config.use_amp,
#                 task=single_task,
#                 robot_type=robot.robot_type,
#             )
#             action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
#         elif policy is None and isinstance(teleop, Teleoperator):
#             action = teleop.get_action()
#         elif policy is None and isinstance(teleop, list):
#             # TODO(pepijn, steven): clean the record loop for use of multiple robots (possibly with pipeline)
#             arm_action = teleop_arm.get_action()
#             arm_action = {f"arm_{k}": v for k, v in arm_action.items()}

#             keyboard_action = teleop_keyboard.get_action()
#             base_action = robot._from_keyboard_to_base_action(keyboard_action)

#             action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
#         else:
#             logging.info(
#                 "No policy or teleoperator provided, skipping action generation."
#                 "This is likely to happen when resetting the environment without a teleop device."
#                 "The robot won't be at its rest position at the start of the next episode."
#             )
#             continue

#         # Action can eventually be clipped using `max_relative_target`,
#         # so action actually sent is saved in the dataset.
#         sent_action = robot.send_action(action)

#         if dataset is not None:
#             action_frame = build_dataset_frame(dataset.features, sent_action, prefix="action")
#             frame = {**observation_frame, **action_frame}
#             dataset.add_frame(frame, task=single_task)

#         if display_data:
#             log_rerun_data(observation, action)

#         dt_s = time.perf_counter() - start_loop_t
#         busy_wait(1 / fps - dt_s)

#         timestamp = time.perf_counter() - start_episode_t


# @parser.wrap()
# def record(cfg: RecordConfig) -> LeRobotDataset:
#     init_logging()

#     # 检查是否提供了 single_task，如果没有，提示用户输入
#     if cfg.dataset.single_task is None:
#         print("请输入任务描述（例如：'Place the black pen inside the pen holder'）：")
#         user_task = input().strip()
#         if not user_task:
#             raise ValueError("任务描述不能为空！请输入有效的任务描述。")
#         cfg.dataset.single_task = user_task
#         logging.info(f"用户输入的任务描述: {cfg.dataset.single_task}")
#     elif not cfg.dataset.single_task:
#         raise ValueError("任务描述不能为空！请输入有效的任务描述。")

#     logging.info(pformat(asdict(cfg)))
#     if cfg.display_data:
#         _init_rerun(session_name="recording")

#     robot = make_robot_from_config(cfg.robot)
#     teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None

#     action_features = hw_to_dataset_features(robot.action_features, "action", cfg.dataset.video)
#     obs_features = hw_to_dataset_features(robot.observation_features, "observation", cfg.dataset.video)
#     dataset_features = {**action_features, **obs_features}

#     if cfg.resume:
#         dataset = LeRobotDataset(
#             cfg.dataset.repo_id,
#             root=cfg.dataset.root,
#         )
#         if hasattr(robot, "cameras") and len(robot.cameras) > 0:
#             dataset.start_image_writer(
#                 num_processes=cfg.dataset.num_image_writer_processes,
#                 num_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras),
#             )
#         sanity_check_dataset_robot_compatibility(dataset, robot, cfg.dataset.fps, dataset_features)
#     else:
#         sanity_check_dataset_name(cfg.dataset.repo_id, cfg.policy)
#         dataset = LeRobotDataset.create(
#             cfg.dataset.repo_id,
#             cfg.dataset.fps,
#             root=cfg.dataset.root,
#             robot_type=robot.name,
#             features=dataset_features,
#             use_videos=cfg.dataset.video,
#             image_writer_processes=cfg.dataset.num_image_writer_processes,
#             image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras),
#         )

#     policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)

#     robot.connect()
#     if teleop is not None:
#         teleop.connect()

#     listener, events = init_keyboard_listener()

#     recorded_episodes = 0
#     while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
#         log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
#         record_loop(
#             robot=robot,
#             events=events,
#             fps=cfg.dataset.fps,
#             teleop=teleop,
#             policy=policy,
#             dataset=dataset,
#             control_time_s=cfg.dataset.episode_time_s,
#             single_task=cfg.dataset.single_task,
#             display_data=cfg.display_data,
#         )

#         if not events["stop_recording"] and (
#             (recorded_episodes < cfg.dataset.num_episodes - 1) or events["rerecord_episode"]
#         ):
#             log_say("Reset the environment", cfg.play_sounds)
#             record_loop(
#                 robot=robot,
#                 events=events,
#                 fps=cfg.dataset.fps,
#                 teleop=teleop,
#                 control_time_s=cfg.dataset.reset_time_s,
#                 single_task=cfg.dataset.single_task,
#                 display_data=cfg.display_data,
#             )

#         if events["rerecord_episode"]:
#             log_say("Re-record episode", cfg.play_sounds)
#             events["rerecord_episode"] = False
#             events["exit_early"] = False
#             dataset.clear_episode_buffer()
#             continue

#         dataset.save_episode()
#         recorded_episodes += 1

#     log_say("Stop recording", cfg.play_sounds, blocking=True)

#     robot.disconnect()
#     if teleop is not None:
#         teleop.disconnect()

#     if not is_headless() and listener is not None:
#         listener.stop()

#     if cfg.dataset.push_to_hub:
#         dataset.push_to_hub(tags=cfg.dataset.tags, private=cfg.dataset.private)

#     log_say("Exiting", cfg.play_sounds)
#     return dataset
# if __name__ == "__main__":
#     record()











####################################################  任务二：只输入任务描述，自动生成任务描述的指令  ############################################
"""
优化：参数已硬编码到脚本中，运行时直接提示用户输入任务描述。

    在 DatasetRecordConfig 类中，将命令行参数硬编码为默认值：

    在 RecordConfig 类中，硬编码机器人和策略配置
"""
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# """
# wwq_record.py 是 LeRobot 框架中用于录制机器人数据集的脚本，主要功能包括：

#     数据采集：从机器人（例如 so100_follower）和摄像头（例如 OpenCV 摄像头）获取观测数据（图像、状态）。
#     动作生成：通过远程操作（teleop）或策略（policy）生成机器人动作。
#     数据集保存：将观测和动作保存为数据集，存储到本地（--dataset.root）或推送到 Hugging Face Hub（--dataset.push_to_hub）。
#     实时控制：以指定帧率（--dataset.fps）控制机器人动作和数据录制。

# 优化：参数已硬编码到脚本中，运行时直接提示用户输入任务描述。
# """

# import logging
# import time
# from dataclasses import asdict, dataclass
# from pathlib import Path
# from pprint import pformat
# from typing import List

# from lerobot.cameras import CameraConfig
# from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
# from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig
# from lerobot.configs.policies import PreTrainedConfig
# from lerobot.datasets.image_writer import safe_stop_image_writer
# from lerobot.datasets.lerobot_dataset import LeRobotDataset
# from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
# from lerobot.policies.factory import make_policy
# from lerobot.policies.pretrained import PreTrainedPolicy
# from lerobot.robots import Robot, RobotConfig, koch_follower, make_robot_from_config, so100_follower, so101_follower
# from lerobot.teleoperators import Teleoperator, TeleoperatorConfig, koch_leader, make_teleoperator_from_config, so100_leader, so101_leader
# from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop
# from lerobot.utils.control_utils import init_keyboard_listener, is_headless, predict_action, sanity_check_dataset_name, sanity_check_dataset_robot_compatibility
# from lerobot.utils.robot_utils import busy_wait
# from lerobot.utils.utils import get_safe_torch_device, init_logging, log_say
# from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
# from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
# # 定义数据集录制的参数
# @dataclass
# class DatasetRecordConfig:
#     repo_id: str = "datasets/eval_wwq_pencil"
#     single_task: str | None = None
#     root: str | Path = "/home/zhq2004/XArm/lerobot_smolvla/datasets/eval_pick_place_red_smolvla2"
#     fps: int = 30
#     episode_time_s: int | float = 30
#     reset_time_s: int | float = 60
#     num_episodes: int = 1
#     video: bool = True
#     push_to_hub: bool = False
#     private: bool = False
#     tags: list[str] | None = None
#     num_image_writer_processes: int = 0
#     num_image_writer_threads_per_camera: int = 4

# # 定义整体录制配置，包括机器人、数据集和控制方式


# @dataclass
# class RecordConfig:
#     robot: RobotConfig = SO100FollowerConfig(
#         id="my_awesome_follower_arm",
#         port="/dev/ttyACM14",
#         cameras={
#             "top": OpenCVCameraConfig(index_or_path=2, width=640, height=480, fps=30),
#             "laptop": OpenCVCameraConfig(index_or_path=4, width=640, height=480, fps=30),
#         }
#     )
#     dataset: DatasetRecordConfig = DatasetRecordConfig()
#     teleop: TeleoperatorConfig | None = None
#     policy: PreTrainedConfig = PreTrainedConfig.from_pretrained(
#         "/home/zhq2004/XArm/lerobot_smolvla/outputs/train/smolvla_so100_pencil/checkpoints/160000/pretrained_model"
#     )
#     display_data: bool = False
#     play_sounds: bool = True
#     resume: bool = False

#     def __post_init__(self):
#         if self.teleop is None and self.policy is None:
#             raise ValueError("Choose a policy, a teleoperator or both to control the robot")

# @safe_stop_image_writer
# def record_loop(
#     robot: Robot,
#     events: dict,
#     fps: int,
#     dataset: LeRobotDataset | None = None,
#     teleop: Teleoperator | List[Teleoperator] | None = None,
#     policy: PreTrainedPolicy | None = None,
#     control_time_s: int | None = None,
#     single_task: str | None = None,
#     display_data: bool = False,
# ):
#     if dataset is not None and dataset.fps != fps:
#         raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).")

#     teleop_arm = teleop_keyboard = None
#     if isinstance(teleop, list):
#         teleop_keyboard = next((t for t in teleop if isinstance(t, KeyboardTeleop)), None)
#         teleop_arm = next(
#             (
#                 t
#                 for t in teleop
#                 if isinstance(t, (so100_leader.SO100Leader, so101_leader.SO101Leader, koch_leader.KochLeader))
#             ),
#             None,
#         )

#         if not (teleop_arm and teleop_keyboard and len(teleop) == 2 and robot.name == "lekiwi_client"):
#             raise ValueError(
#                 "For multi-teleop, the list must contain exactly one KeyboardTeleop and one arm teleoperator. Currently only supported for LeKiwi robot."
#             )

#     if policy is not None:
#         policy.reset()

#     timestamp = 0
#     start_episode_t = time.perf_counter()
#     while timestamp < control_time_s:
#         start_loop_t = time.perf_counter()

#         if events["exit_early"]:
#             events["exit_early"] = False
#             break

#         observation = robot.get_observation()

#         if policy is not None or dataset is not None:
#             observation_frame = build_dataset_frame(dataset.features, observation, prefix="observation")

#         if policy is not None:
#             action_values = predict_action(
#                 observation_frame,
#                 policy,
#                 get_safe_torch_device(policy.config.device),
#                 policy.config.use_amp,
#                 task=single_task,
#                 robot_type=robot.robot_type,
#             )
#             action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
#         elif policy is None and isinstance(teleop, Teleoperator):
#             action = teleop.get_action()
#         elif policy is None and isinstance(teleop, list):
#             arm_action = teleop_arm.get_action()
#             arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
#             keyboard_action = teleop_keyboard.get_action()
#             base_action = robot._from_keyboard_to_base_action(keyboard_action)
#             action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
#         else:
#             logging.info(
#                 "No policy or teleoperator provided, skipping action generation."
#                 "This is likely to happen when resetting the environment without a teleop device."
#                 "The robot won't be at its rest position at the start of the next episode."
#             )
#             continue

#         sent_action = robot.send_action(action)

#         if dataset is not None:
#             action_frame = build_dataset_frame(dataset.features, sent_action, prefix="action")
#             frame = {**observation_frame, **action_frame}
#             dataset.add_frame(frame, task=single_task)

#         if display_data:
#             log_rerun_data(observation, action)

#         dt_s = time.perf_counter() - start_loop_t
#         busy_wait(1 / fps - dt_s)

#         timestamp = time.perf_counter() - start_episode_t

# def record() -> LeRobotDataset:
#     init_logging()

#     # 使用硬编码的配置
#     cfg = RecordConfig()

#     # 强制提示用户输入任务描述
#     print("请输入任务描述（例如：'Place the black pen inside the pen holder'）：")
#     user_task = input().strip()
#     if not user_task:
#         raise ValueError("任务描述不能为空！请输入有效的任务描述。")
#     cfg.dataset.single_task = user_task
#     logging.info(f"用户输入的任务描述: {cfg.dataset.single_task}")

#     logging.info(pformat(asdict(cfg)))
#     if cfg.display_data:
#         _init_rerun(session_name="recording")

#     robot = make_robot_from_config(cfg.robot)
#     teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None

#     action_features = hw_to_dataset_features(robot.action_features, "action", cfg.dataset.video)
#     obs_features = hw_to_dataset_features(robot.observation_features, "observation", cfg.dataset.video)
#     dataset_features = {**action_features, **obs_features}

#     if cfg.resume:
#         dataset = LeRobotDataset(
#             cfg.dataset.repo_id,
#             root=cfg.dataset.root,
#         )
#         if hasattr(robot, "cameras") and len(robot.cameras) > 0:
#             dataset.start_image_writer(
#                 num_processes=cfg.dataset.num_image_writer_processes,
#                 num_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras),
#             )
#         sanity_check_dataset_robot_compatibility(dataset, robot, cfg.dataset.fps, dataset_features)
#     else:
#         sanity_check_dataset_name(cfg.dataset.repo_id, cfg.policy)
#         dataset = LeRobotDataset.create(
#             cfg.dataset.repo_id,
#             cfg.dataset.fps,
#             root=cfg.dataset.root,
#             robot_type=robot.name,
#             features=dataset_features,
#             use_videos=cfg.dataset.video,
#             image_writer_processes=cfg.dataset.num_image_writer_processes,
#             image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras),
#         )

#     policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)

#     robot.connect()
#     if teleop is not None:
#         teleop.connect()

#     listener, events = init_keyboard_listener()

#     recorded_episodes = 0
#     while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
#         log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
#         record_loop(
#             robot=robot,
#             events=events,
#             fps=cfg.dataset.fps,
#             teleop=teleop,
#             policy=policy,
#             dataset=dataset,
#             control_time_s=cfg.dataset.episode_time_s,
#             single_task=cfg.dataset.single_task,
#             display_data=cfg.display_data,
#         )

#         if not events["stop_recording"] and (
#             (recorded_episodes < cfg.dataset.num_episodes - 1) or events["rerecord_episode"]
#         ):
#             log_say("Reset the environment", cfg.play_sounds)
#             record_loop(
#                 robot=robot,
#                 events=events,
#                 fps=cfg.dataset.fps,
#                 teleop=teleop,
#                 control_time_s=cfg.dataset.reset_time_s,
#                 single_task=cfg.dataset.single_task,
#                 display_data=cfg.display_data,
#             )

#         if events["rerecord_episode"]:
#             log_say("Re-record episode", cfg.play_sounds)
#             events["rerecord_episode"] = False
#             events["exit_early"] = False
#             dataset.clear_episode_buffer()
#             continue

#         dataset.save_episode()
#         recorded_episodes += 1

#     log_say("Stop recording", cfg.play_sounds, blocking=True)

#     robot.disconnect()
#     if teleop is not None:
#         teleop.disconnect()

#     if not is_headless() and listener is not None:
#         listener.stop()

#     if cfg.dataset.push_to_hub:
#         dataset.push_to_hub(tags=cfg.dataset.tags, private=cfg.dataset.private)

#     log_say("Exiting", cfg.play_sounds)
#     return dataset

# if __name__ == "__main__":
#     record()


